package com.xiashitech.agent.instrumentation.template;

import com.xiashitech.agent.instrumentation.base.DebugUtils;
import com.xiashitech.agent.instrumentation.base.InstProcessor;
import com.xiashitech.agent.instrumentation.base.context.TraceContextVars;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.Baggage;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.common.Attributes;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.common.AttributesBuilder;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;

/**
 * 插桩的业务逻辑处理器
 */
public class TemplateProcessor extends InstProcessor<TemplateRequestData, Object> {

    public TemplateProcessor(String spanName) {
        super(spanName);
    }

    /**
     * 业务逻辑的onStart处理，可以执行获取baggage信息，设置attribute等操作
     */
    @Override
    public void onXSStart(AttributesBuilder attributesBuilder, Context context, TemplateRequestData requestData) {
        //从context中获取baggage
        Baggage baggage = Baggage.fromContext(context);

        try {
            //start节点的业务处理
        } catch (Exception e) {
            DebugUtils.printError4Debug(e, attributesBuilder); //打印异常信息，并输入到attributes，方便zipkin中查看
        }

        //设置链路追踪attributes
        attributesBuilder.put("xs.key1", requestData.getData1()).put("xs.key1", requestData.getData2());
    }

    /**
     * 在start阶段会调用该方法，更新baggage信息的操作应放在此处执行
     * @return 新的context对象，如果返回null则不更新context
     */
    @Override
    public Context setBaggageAndContext(Context context, TemplateRequestData requestData, Attributes attributes) {
        //更新baggage信息和context对象
        context = updateBaggage(context, builder ->
                builder.put("xs.key1", requestData.getData1()).put("xs.key2", requestData.getData2()));

        //返回新的context
        return context;
    }

    /**
     * 业务逻辑的onXSEnd处理，可执行设置TraceContextVars的操作，以实现向后续节点传递数据
     */
    @Override
    public void onXSEnd(AttributesBuilder attributesBuilder, Context context, TemplateRequestData requestData, Object responseData, Throwable throwable) {
        try {
            //处理responseData
            //end节点的业务处理
        } catch (Exception e) {
            DebugUtils.printError4Debug(e, attributesBuilder); //打印异常信息，并输入到attributes，方便zipkin中查看
        }

        //设置链路追踪attributes
        attributesBuilder.put("xs.key3", "value3");

        //设置TraceContextVars，向后传递数据
        TraceContextVars.get().set("xs.key3", "value3");
    }

}
